package com.mini.rpc.provider.registry.loadbalancer;

import com.mini.rpc.common.ServiceMeta;
import com.mini.rpc.provider.registry.exception.NoAvailableProviderException;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * Hash一致性, 负载均衡算法
 * <p>
 * 尽可能保证每个节点分摊到的流量尽可能均衡
 */
public class ZKConsistentHashLoadBalancer implements ServiceLoadBalancer<ServiceMeta> {
    /**
     * 虚拟节点数量
     */
    private final static int VIRTUAL_NODE_SIZE = 10;
    private final static String VIRTUAL_NODE_SPLIT = "#";
    /**
     * 定义treeMap
     */
    private TreeMap<Integer, ServiceMeta> ring = new TreeMap<>();

    /**
     * Hash一致性负载均衡算法获取一个节点
     *
     * @param servers  服务节点列表
     * @param hashCode 客服端节点的hashCode
     * @return 服务节点实例
     */
    @Override
    public ServiceMeta select(List<ServiceMeta> servers, int hashCode) {
        if (servers.size() == 0) {
            throw new NoAvailableProviderException("no available provider");
        }
        // 构建hash环
        TreeMap<Integer, ServiceMeta> ring = makeConsistentHashRing(servers);
        // 根据hashCode分配结点
        return allocateNode(ring, hashCode);
    }

    private ServiceMeta allocateNode(TreeMap<Integer, ServiceMeta> ring, int hashCode) {
        // 顺时针找到第一个节点
        Map.Entry<Integer, ServiceMeta> entry = ring.ceilingEntry(hashCode);
        if (entry == null) {
            // 如果没有大于hashCode的节点, 直接取第一个
            entry = ring.firstEntry();
        }
        return entry.getValue();
    }

    private TreeMap<Integer, ServiceMeta> makeConsistentHashRing(List<ServiceMeta> servers) {
        // TreeMap<Integer, ServiceMeta> ring = new TreeMap<>();
        for (ServiceMeta instance : servers) {
            for (int i = 0; i < VIRTUAL_NODE_SIZE; i++) {
                ring.put((buildServiceInstanceKey(instance) + VIRTUAL_NODE_SPLIT + i).hashCode(), instance);
            }
        }
        return ring;
    }

    private String buildServiceInstanceKey(ServiceMeta instance) {
        return String.join(":", instance.getServiceAddr(), String.valueOf(instance.getServicePort()));
    }

}